Skip to content

feat: Support QLoRA fine-tuning in post-training#3702

Draft
RexBearIU wants to merge 9 commits into
mainfrom
jackyf/qlora
Draft

feat: Support QLoRA fine-tuning in post-training#3702
RexBearIU wants to merge 9 commits into
mainfrom
jackyf/qlora

Conversation

@RexBearIU
Copy link
Copy Markdown
Collaborator

@RexBearIU RexBearIU commented Apr 20, 2026 •
edited
Loading

Description

This PR extends our existing LoRA integration to support Quantized LoRA (QLoRA) using the Qwix library for our NNX-based models. It adds configuration options for quantization types and tile sizes, along with necessary logic
fixes to correctly handle quantized node metadata and parameter traversal within the NNX framework.

Key Changes

  • Config Additions: Added lora_weight_qtype (e.g., nf4) and lora_tile_size to sft.yml and the LoRA type class to enable QLoRA configurations.
  • NNX Decoder Enhancements:
    • Added metadata preservation (stash_origin_metadata and restore_origin_metadata) to correctly handle partitioning specs across scan boundaries in NNXDecoder.
    • Introduced fix_node_rank logic to dynamically adjust PartitionSpec rank constraints to match parameter shapes during scanning.
  • Qwix Provider Patching (lora_utils.py):
    • Automatically switches between LoRA and QLoRA providers based on the presence of lora_weight_qtype.
    • Implemented _patch_qwix_for_maxtext to resolve integration issues:
      • PTQ Patch: Intercepts jax.numpy.asarray to correctly handle nnx.State arrays wrapped as ptq.QArray objects.
      • Parameter Traversal Patch (find_param): Replaces flax_util.find_param with a custom implementation that accurately searches nnx.Module trees and jax.core.Tracer graphs to find the correct node references for
        quantization.

Tests

Updated lora_utils_test.py to cover the new optional QLoRA config fields.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@RexBearIU RexBearIU closed this Apr 20, 2026
@RexBearIU RexBearIU reopened this Apr 21, 2026
@RexBearIU RexBearIU force-pushed the jackyf/feat/lora-nnx branch 3 times, most recently from 669b501 to 9634d50 Compare April 21, 2026 08:48
@RexBearIU RexBearIU force-pushed the jackyf/qlora branch 2 times, most recently from a2739ab to fc27813 Compare April 22, 2026 08:01
@RexBearIU RexBearIU force-pushed the jackyf/feat/lora-nnx branch 22 times, most recently from 8531d67 to de940c3 Compare April 29, 2026 12:28
@RexBearIU RexBearIU force-pushed the jackyf/qlora branch 8 times, most recently from 2d237f8 to f8a7f1b Compare May 4, 2026 11:06
@RexBearIU RexBearIU force-pushed the jackyf/feat/lora-nnx branch from 5a7bedb to 397f319 Compare May 5, 2026 08:03
@codecov
Copy link
Copy Markdown

codecov Bot commented May 5, 2026 •
edited
Loading

Codecov Report

❌ Patch coverage is 69.69697% with 20 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/utils/maxtext_utils_nnx.py 66.07% 8 Missing and 11 partials ⚠️
src/maxtext/layers/nnx_decoders.py 87.50% 0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@RexBearIU RexBearIU changed the title Jackyf/qlora feat: Support QLoRA fine-tuning in post-training May 5, 2026
@RexBearIU RexBearIU force-pushed the jackyf/feat/lora-nnx branch 5 times, most recently from 391664f to 1eb5953 Compare May 8, 2026 11:11
@shralex shralex force-pushed the jackyf/feat/lora-nnx branch 2 times, most recently from 3954ce8 to 6face5b Compare May 8, 2026 15:01
@RexBearIU RexBearIU force-pushed the jackyf/feat/lora-nnx branch from 6face5b to eced9d7 Compare May 11, 2026 10:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jacoguzo jacoguzo Awaiting requested review from jacoguzo jacoguzo will be requested when the pull request is marked ready for review jacoguzo is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon will be requested when the pull request is marked ready for review bvandermoon is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC will be requested when the pull request is marked ready for review SurbhiJainUSC is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 will be requested when the pull request is marked ready for review richjames0 is a code owner

@shralex shralex Awaiting requested review from shralex shralex will be requested when the pull request is marked ready for review shralex is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk will be requested when the pull request is marked ready for review gobbleturk is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan will be requested when the pull request is marked ready for review RissyRan is a code owner

@gagika gagika Awaiting requested review from gagika gagika will be requested when the pull request is marked ready for review gagika is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha will be requested when the pull request is marked ready for review A9isha is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 will be requested when the pull request is marked ready for review jiangjy1982 is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla will be requested when the pull request is marked ready for review vipannalla is a code owner

@parambole parambole Awaiting requested review from parambole parambole will be requested when the pull request is marked ready for review parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin will be requested when the pull request is marked ready for review shuningjin is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 will be requested when the pull request is marked ready for review suexu1025 is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande will be requested when the pull request is marked ready for review NicoGrande is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google will be requested when the pull request is marked ready for review jesselu-google is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng will be requested when the pull request is marked ready for review NuojCheng is a code owner

@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson will be requested when the pull request is marked ready for review abhinavclemson is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit will be requested when the pull request is marked ready for review khatwanimohit is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo will be requested when the pull request is marked ready for review hengtaoguo is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei will be requested when the pull request is marked ready for review aireenmei is a code owner

@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 will be requested when the pull request is marked ready for review dipannita08 is a code owner

@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git will be requested when the pull request is marked ready for review igorts-git is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@RexBearIU